# -*- coding: utf-8 -*-
# train.py

# from scipy.sparse import csr_matrix

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import GridSearchCV
# from sklearn.feature_extraction.text import TfidfTransformer
import pickle
import sys
import numpy as np

from sklearn import metrics

import matplotlib.pyplot as plt

from sklearn.metrics import auc


def get_confusion_matrix(target, pred):
    return confusion_matrix(target, pred)


if __name__ == '__main__':
    data_file = sys.argv[1]
    lg_file = sys.argv[2]

    X, target = pickle.load(open(data_file, 'rb'))

    train_x, test_x, train_y, test_y = train_test_split(X, target, test_size=0.20, random_state=30)
    # logreg = LogisticRegression(penalty='l1', C=1e2, max_iter=100, multi_class='ovr')
    logreg = LogisticRegression(penalty='l1', max_iter=100, multi_class='ovr')

    # 使用cross validation 选择最佳正则化参数
    params = {"C": [1e-3, 1e-2, 1e-1, 1, 1e1, 1e2]}
    cv = StratifiedKFold(n_splits=10, shuffle=True)
    logreg_cv = GridSearchCV(logreg, params, cv=cv, n_jobs=6)

    logreg_cv.fit(train_x, train_y)

    best_logreg = logreg_cv.best_estimator_
    print(logreg_cv.cv_results_)

    # 保存训练好的模型
    with open(lg_file, 'wb+') as fd:
        pickle.dump(best_logreg, fd)

    # for k, v in zip(best_logreg.predict(test_x[0:100]), test_y[0:100]):
    #     print("%s %s" % (k,v))

    print(best_logreg)

    pred_y = np.array(best_logreg.predict(test_x))
    # print(err_cnt)
    # print(err_cnt/len(test_y))

    from sklearn.metrics import classification_report, confusion_matrix
    print(classification_report(test_y, pred_y))
    # conf_mat = confusion_matrix(test_y, pred_y)

    model = best_logreg
    print('class,' + ','.join(model.classes_))
    conf_mat = get_confusion_matrix([i.strip() for i in test_y], [i.strip() for i in pred_y])
    shape = conf_mat.shape
    k = shape[0]

    tp = [conf_mat[i][i] for i in range(k)]

    truth = np.sum(conf_mat, 1)
    predict = np.sum(conf_mat, 0)
    for e, t in enumerate(conf_mat):
        c = model.classes_[e]
        print(c + ',' + ','.join([str(i) for i in t]))

    print('tp,' + ','.join([str(t) for t in tp]))
    print('truth,' + ','.join([str(t) for t in truth]))
    print('predict,' + ','.join([str(t) for t in predict]))

    # plot roc
    if len(model.classes_ == 2):
        pass


